#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
比亚迪股票数据分析与可视化工具

本脚本用于分析比亚迪股票历史行情数据，包括数据读取、描述性分析、K线图绘制和相关性分析等功能。
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import os
import argparse
from typing import Dict, List, Tuple, Optional

# 设置中文字体，确保中文正常显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False


class BYDStockAnalyzer:
    """
    比亚迪股票数据分析类
    用于读取和分析比亚迪股票历史行情数据
    """
    
    def __init__(self, data_file: str):
        """
        初始化BYDStockAnalyzer类
        
        Args:
            data_file: 比亚迪股票数据文件路径
        """
        self.data_file = data_file
        self.data = None
        self.analysis_results = {}
        
    def load_data(self) -> bool:
        """
        加载比亚迪股票数据
        
        Returns:
            bool: 数据加载是否成功
        """
        try:
            # 读取Excel文件
            self.data = pd.read_excel(self.data_file)
            
            # 检查数据结构
            # 检查数据结构
            if '收盘' not in self.data.columns:
                # 如果没有收盘列，尝试设置第一列为日期
                self.data.columns = ['日期', '开盘', '最高', '最低', '收盘', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
            
            # 将日期列转换为datetime类型
            self.data['日期'] = pd.to_datetime(self.data['日期'])
            
            # 设置日期为索引
            self.data.set_index('日期', inplace=True)
            
            print(f"成功加载数据，共{len(self.data)}行记录")
            print(f"数据时间范围：{self.data.index.min()} 至 {self.data.index.max()}")
            return True
        except Exception as e:
            print(f"数据加载失败: {e}")
            return False
    
    def get_data_preview(self, n: int = 5) -> pd.DataFrame:
        """
        获取数据前n行预览
        
        Args:
            n: 预览行数
        
        Returns:
            pd.DataFrame: 数据前n行
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        return self.data.head(n)
    
    def descriptive_analysis(self) -> Dict:
        """
        对数据进行描述性统计分析
        
        Returns:
            Dict: 描述性统计结果
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        
        # 计算基本统计量
        basic_stats = self.data.describe()
        
        # 初始化变量
        volatility = None
        annual_volatility = None
        max_gain = None
        max_loss = None
        win_rate = None
        
        # 计算涨跌幅
        if '收盘' in self.data.columns:
            self.data['涨跌幅'] = self.data['收盘'].pct_change() * 100
            
            # 计算波动率（使用涨跌幅的标准差）
            volatility = self.data['涨跌幅'].std()
            annual_volatility = volatility * np.sqrt(252)  # 年化波动率
            
            # 计算最大涨幅和最大跌幅
            max_gain = self.data['涨跌幅'].max()
            max_loss = self.data['涨跌幅'].min()
            
            # 计算胜率（上涨天数占比）
            win_rate = len(self.data[self.data['涨跌幅'] > 0]) / len(self.data[self.data['涨跌幅'].notna()])
        
        # 保存分析结果
        self.analysis_results = {
            'basic_stats': basic_stats,
            'volatility': volatility,
            'annual_volatility': annual_volatility,
            'max_gain': max_gain,
            'max_loss': max_loss,
            'win_rate': win_rate
        }
        
        return self.analysis_results
    
    def plot_k_line(self, start_date: Optional[str] = None, end_date: Optional[str] = None, save_path: Optional[str] = None):
        """
        绘制K线图
        
        Args:
            start_date: 开始日期，格式为'YYYY-MM-DD'
            end_date: 结束日期，格式为'YYYY-MM-DD'
            save_path: 图表保存路径，None表示不保存
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        
        # 筛选时间范围
        plot_data = self.data.copy()
        if start_date:
            plot_data = plot_data[plot_data.index >= pd.to_datetime(start_date)]
        if end_date:
            plot_data = plot_data[plot_data.index <= pd.to_datetime(end_date)]
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(15, 8))
        
        # 获取日期索引的数值表示，用于正确绘制K线
        date_positions = range(len(plot_data))
        dates = plot_data.index
        
        # 绘制蜡烛图的实体部分
        for i, (date, row) in enumerate(plot_data.iterrows()):
            if row['收盘'] >= row['开盘']:
                color = 'red'  # 上涨为红色
                ax.bar(i, row['收盘'] - row['开盘'], bottom=row['开盘'], width=0.6, color=color)
            else:
                color = 'green'  # 下跌为绿色
                ax.bar(i, row['开盘'] - row['收盘'], bottom=row['收盘'], width=0.6, color=color)
            
            # 绘制上下影线
            ax.plot([i, i], [row['最低'], row['最高']], color=color, linewidth=1)
        
        # 设置日期标签
        ax.set_xticks(date_positions[::10])  # 每10个交易日显示一个日期
        ax.set_xticklabels([date.strftime('%Y-%m-%d') for date in dates[::10]])
        
        # 设置标题和标签
        ax.set_title('比亚迪股票K线图')
        ax.set_xlabel('日期')
        ax.set_ylabel('价格')
        
        # 设置网格线
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # 自动旋转日期标签
        plt.xticks(rotation=45)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        if save_path:
            plt.savefig(save_path, dpi=300)
            print(f"K线图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def plot_price_trend(self, save_path: Optional[str] = None):
        """
        绘制价格趋势图
        
        Args:
            save_path: 图表保存路径，None表示不保存
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(15, 8))
        
        # 绘制收盘价趋势
        ax.plot(self.data.index, self.data['收盘'], label='收盘', linewidth=2)
        
        # 设置标题和标签
        ax.set_title('比亚迪股票价格趋势')
        ax.set_xlabel('日期')
        ax.set_ylabel('价格')
        
        # 添加图例
        ax.legend()
        
        # 设置网格线
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # 自动旋转日期标签
        plt.xticks(rotation=45)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        if save_path:
            plt.savefig(save_path, dpi=300)
            print(f"价格趋势图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def plot_volume(self, save_path: Optional[str] = None):
        """
        绘制成交量图
        
        Args:
            save_path: 图表保存路径，None表示不保存
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(15, 6))
        
        # 绘制成交量柱状图
        ax.bar(self.data.index, self.data['成交量'], label='成交量', color='blue', alpha=0.7)
        
        # 设置标题和标签
        ax.set_title('比亚迪股票成交量')
        ax.set_xlabel('日期')
        ax.set_ylabel('成交量')
        
        # 添加图例
        ax.legend()
        
        # 设置网格线
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # 自动旋转日期标签
        plt.xticks(rotation=45)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        if save_path:
            plt.savefig(save_path, dpi=300)
            print(f"成交量图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def correlation_analysis(self, save_path: Optional[str] = None):
        """
        计算相关性并绘制热力图
        
        Args:
            save_path: 图表保存路径，None表示不保存
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        
        # 计算相关性矩阵
        corr_matrix = self.data.corr()
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 绘制热力图
        sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', linewidths=0.5, ax=ax)
        
        # 设置标题
        ax.set_title('比亚迪股票数据相关性热力图')
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        if save_path:
            plt.savefig(save_path, dpi=300)
            print(f"相关性热力图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return corr_matrix
    
    def plot_2023_apr_jun_kline(self, save_path: Optional[str] = None):
        """
        绘制2023年4月至2023年6月的K线图
        
        Args:
            save_path: 图表保存路径，None表示不保存
        """
        self.plot_k_line(start_date='2023-04-01', end_date='2023-06-30', save_path=save_path)
    
    def plot_2023_year_chart(self, save_path: Optional[str] = None):
        """
        绘制2023年全年的收盘价时序图
        
        Args:
            save_path: 图表保存路径，None表示不保存
        """
        # 筛选2023年数据
        data_2023 = self.data[self.data.index.year == 2023]
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(15, 8))
        
        # 绘制收盘价时序图
        ax.plot(data_2023.index, data_2023['收盘'], label='2023年收盘', linewidth=2)
        
        # 设置标题和标签
        ax.set_title('比亚迪股票2023年收盘价时序图')
        ax.set_xlabel('日期')
        ax.set_ylabel('收盘')
        
        # 添加图例
        ax.legend()
        
        # 设置网格线
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # 自动旋转日期标签
        plt.xticks(rotation=45)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        if save_path:
            plt.savefig(save_path, dpi=300)
            print(f"2023年收盘价时序图已保存至: {save_path}")
        
        # 显示图表
        plt.show()


def main():
    """
    主函数，用于演示比亚迪股票数据分析流程
    """
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='比亚迪股票数据分析工具')
    parser.add_argument('--data_file', 
                        default='c:\\Users\\YEDX\\financial-big-data-course\\学习通第二次作业\\比亚迪后复权历史行情数据.xlsx',
                        help='比亚迪股票数据文件路径')
    parser.add_argument('--output_dir',
                        default='c:\\Users\\YEDX\\financial-big-data-course\\chapter4\\byd_stock_analysis\\output',
                        help='图表保存目录')
    parser.add_argument('--skip_plots', action='store_true', help='跳过图表绘制')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 创建分析器实例
    analyzer = BYDStockAnalyzer(args.data_file)
    
    # 加载数据
    if not analyzer.load_data():
        print("程序终止")
        return
    
    # 显示数据前5行
    print("\n数据预览:")
    print(analyzer.get_data_preview())
    
    # 进行描述性分析
    print("\n进行描述性分析...")
    analysis_results = analyzer.descriptive_analysis()
    print("基本统计量:")
    print(analysis_results['basic_stats'])
    
    if 'volatility' in analysis_results and analysis_results['volatility'] is not None:
        print(f"\n日波动率: {analysis_results['volatility']:.2f}%")
        print(f"年化波动率: {analysis_results['annual_volatility']:.2f}%")
        print(f"最大涨幅: {analysis_results['max_gain']:.2f}%")
        print(f"最大跌幅: {analysis_results['max_loss']:.2f}%")
        print(f"胜率: {analysis_results['win_rate']:.2%}")
    
    # 创建保存图表的目录
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 如果跳过图表绘制，则直接返回
    if args.skip_plots:
        print("\n已跳过图表绘制，分析完成！")
        return
    
    try:
        # 绘制2023年4月至6月的K线图
        print("\n绘制2023年4月至6月的K线图...")
        analyzer.plot_2023_apr_jun_kline(save_path=os.path.join(output_dir, 'byd_kline_2023_apr_jun.png'))
        
        # 绘制2023年全年收盘价时序图
        print("\n绘制2023年全年收盘价时序图...")
        analyzer.plot_2023_year_chart(save_path=os.path.join(output_dir, 'byd_close_price_2023.png'))
        
        # 绘制价格趋势图
        print("\n绘制价格趋势图...")
        analyzer.plot_price_trend(save_path=os.path.join(output_dir, 'byd_price_trend.png'))
        
        # 绘制成交量图
        print("\n绘制成交量图...")
        analyzer.plot_volume(save_path=os.path.join(output_dir, 'byd_volume.png'))
        
        # 进行相关性分析
        print("\n进行相关性分析...")
        corr_matrix = analyzer.correlation_analysis(save_path=os.path.join(output_dir, 'byd_correlation_heatmap.png'))
        print("相关性矩阵:")
        print(corr_matrix)
        
        print("\n分析完成！")
        print(f"\n所有图表已保存至: {output_dir}")
        print("\n提示：如果你想只进行数据分析而不显示图表，可以使用以下命令：")
        print("python byd_stock_analysis.py --skip_plots")
    except Exception as e:
        print(f"\n绘制图表时发生错误: {e}")
        print("\n数据分析部分已完成，但图表绘制失败。")
        print("你可以尝试使用--skip_plots参数来跳过图表绘制。")


if __name__ == "__main__":
    main()